import os
import argparse

def generate_video_list(input_dir, output_file):
    """
    遍历指定目录下的所有.mp4文件，将其完整路径写入txt文件
    Args:
        input_dir (str): 要遍历的目录路径
        output_file (str): 输出的txt文件路径
    """
    # 确保输入目录存在
    if not os.path.exists(input_dir):
        raise ValueError(f"Input directory {input_dir} does not exist!")
    
    # 获取所有.mp4文件的路径
    video_paths = []
    for root, _, files in os.walk(input_dir):
        for file in files:
            if file.endswith('.mp4'):
                # 获取完整的绝对路径
                full_path = os.path.abspath(os.path.join(root, file))
                # 去掉扩展名
                path_without_ext = os.path.splitext(full_path)[0]
                video_paths.append(path_without_ext)
    
    # 确保找到了视频文件
    if not video_paths:
        print("Warning: No .mp4 files found in the directory!")
        return
    
    print(f"Found {len(video_paths)} video files")
    
    try:
        # 确保输出目录存在
        os.makedirs(os.path.dirname(os.path.abspath(output_file)), exist_ok=True)
        
        # 将路径写入文件
        with open(output_file, 'w', encoding='utf-8') as f:
            for path in video_paths:
                f.write(f"{path}\n")
                
        print(f"Successfully wrote {len(video_paths)} video paths to {output_file}")
        
        # 验证文件是否写入成功
        if os.path.exists(output_file):
            with open(output_file, 'r', encoding='utf-8') as f:
                written_paths = f.readlines()
            print(f"Verified: {len(written_paths)} paths in output file")
            # 打印前几个路径作为示例
            print("\nExample paths:")
            for path in written_paths[:3]:
                print(path.strip())
        else:
            print("Error: Failed to create output file!")
            
    except Exception as e:
        print(f"Error writing to file: {str(e)}")

def main():
    parser = argparse.ArgumentParser(description='Generate a list of video file paths')
    parser.add_argument('--input_dir', type=str, required=True,
                        help='Directory containing the video files')
    parser.add_argument('--output_file', type=str, required=True,
                        help='Output txt file path')
    
    args = parser.parse_args()
    generate_video_list(args.input_dir, args.output_file)

if __name__ == "__main__":
    main()